package cn.springcloud.alibaba.feign.rule;

import cn.springcloud.alibaba.core.constant.Constant;
import com.alibaba.cloud.nacos.ribbon.NacosServer;
import com.google.common.base.Optional;
import com.netflix.loadbalancer.Server;
import com.netflix.loadbalancer.ZoneAvoidanceRule;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * 灰度发布的规则
 */
public class GrayRule extends ZoneAvoidanceRule {

	@Override
	public Server choose(Object key) {
		try {
			//获取所有可用服务
			List<Server> serverList = this.getLoadBalancer().getReachableServers();
			//灰度发布的服务
			List<Server> grayServerList = new ArrayList<>();
			//正常的服务
			List<Server> normalServerList = new ArrayList<>();

			for (Server server : serverList) {
				NacosServer nacosServer = (NacosServer) server;
				if (Boolean.TRUE.toString().equals(nacosServer.getMetadata().get(Constant.GRAY_TAG))) {
					grayServerList.add(server);
				} else {
					normalServerList.add(server);
				}
			}

			//从ThreadLocal中获取灰度标记
			boolean grayTag = GrayRequestContextHolder.getGrayTag();
			//如果被标记为灰度发布，则调用灰度发布的服务
			if (grayTag) {
				return originChoose(grayServerList, key);
			} else {
				return originChoose(normalServerList, key);
			}
		} finally {
			GrayRequestContextHolder.remove();
		}
	}

	public static HttpServletRequest getRequest() {
		return ((ServletRequestAttributes) (Objects.requireNonNull(RequestContextHolder.getRequestAttributes()))).getRequest();
	}

	private Server originChoose(List<Server> noMetaServerList, Object key) {
		Optional<Server> server = getPredicate().chooseRoundRobinAfterFiltering(noMetaServerList, key);
		if (server.isPresent()) {
			return server.get();
		} else {
			return null;
		}
	}
}
